package com.jourwon.spring.boot.core;

import com.baomidou.dynamic.datasource.toolkit.DynamicDataSourceContextHolder;
import org.springframework.dao.DataAccessException;
import org.springframework.jdbc.core.*;
import org.springframework.jdbc.support.KeyHolder;
import org.springframework.jdbc.support.rowset.SqlRowSet;

import java.util.Collection;
import java.util.List;
import java.util.Map;

/**
 * 动态JdbcTemplate
 *
 * @author JourWon
 * @date 2022/3/17
 */
public class DynamicJdbcTemplate extends AbstractDynamicJdbcTemplate {

    private final JdbcTemplate jdbcTemplate;

    public DynamicJdbcTemplate(JdbcTemplate jdbcTemplate,  String primaryDb) {
        super(primaryDb);
        this.jdbcTemplate = jdbcTemplate;
    }

    public <T> T execute(String dsName, ConnectionCallback<T> connectionCallback) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            return this.jdbcTemplate.execute(connectionCallback);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    public <T> T execute(String dsName, StatementCallback<T> statementCallback) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            return this.jdbcTemplate.execute(statementCallback);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    public void execute(String dsName, String sql) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            this.jdbcTemplate.execute(sql);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    public <T> T query(String dsName, String sql, ResultSetExtractor<T> resultSetExtractor) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            return this.jdbcTemplate.query(sql, resultSetExtractor);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    public void query(String dsName, String sql, RowCallbackHandler rowCallbackHandler) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            this.jdbcTemplate.query(sql, rowCallbackHandler);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    public <T> List<T> query(String dsName, String sql, RowMapper<T> rowMapper) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            return this.jdbcTemplate.query(sql, rowMapper);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    public <T> T queryForObject(String dsName, String sql, RowMapper<T> rowMapper) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            return this.jdbcTemplate.queryForObject(sql, rowMapper);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    public <T> T queryForObject(String dsName, String sql, Class<T> aClass) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            return this.jdbcTemplate.queryForObject(sql, aClass);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    public Map<String, Object> queryForMap(String dsName, String sql) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            return this.jdbcTemplate.queryForMap(sql);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    public <T> List<T> queryForList(String dsName, String sql, Class<T> aClass) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            return this.jdbcTemplate.queryForList(sql, aClass);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    public List<Map<String, Object>> queryForList(String dsName, String sql) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            return this.jdbcTemplate.queryForList(sql);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    public SqlRowSet queryForRowSet(String dsName, String sql) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            return this.jdbcTemplate.queryForRowSet(sql);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    public int update(String dsName, String sql) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            return this.jdbcTemplate.update(sql);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    public int[] batchUpdate(String dsName, String... sqls) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            return this.jdbcTemplate.batchUpdate(sqls);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    public <T> T execute(String dsName, PreparedStatementCreator preparedStatementCreator, PreparedStatementCallback<T> preparedStatementCallback) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            return this.jdbcTemplate.execute(preparedStatementCreator, preparedStatementCallback);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    public <T> T execute(String dsName, String sql, PreparedStatementCallback<T> preparedStatementCallback) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            return this.jdbcTemplate.execute(sql, preparedStatementCallback);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    public <T> T query(String dsName, PreparedStatementCreator preparedStatementCreator, ResultSetExtractor<T> resultSetExtractor) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            return this.jdbcTemplate.query(preparedStatementCreator, resultSetExtractor);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    public <T> T query(String dsName, String sql, PreparedStatementSetter preparedStatementSetter, ResultSetExtractor<T> resultSetExtractor) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            return this.jdbcTemplate.query(sql, preparedStatementSetter, resultSetExtractor);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    public <T> T query(String dsName, String sql, Object[] args, int[] argTypes, ResultSetExtractor<T> resultSetExtractor) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            return this.jdbcTemplate.query(sql, args, argTypes, resultSetExtractor);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    public <T> T query(String dsName, String sql, ResultSetExtractor<T> resultSetExtractor, Object... args) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            return this.jdbcTemplate.query(sql, resultSetExtractor, args);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    public void query(String dsName, PreparedStatementCreator preparedStatementCreator, RowCallbackHandler rowCallbackHandler) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            this.jdbcTemplate.query(preparedStatementCreator, rowCallbackHandler);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    public void query(String dsName, String sql, PreparedStatementSetter preparedStatementSetter, RowCallbackHandler rowCallbackHandler) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            this.jdbcTemplate.query(sql, preparedStatementSetter, rowCallbackHandler);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    public void query(String dsName, String sql, Object[] args, int[] argTypes, RowCallbackHandler rowCallbackHandler) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            this.jdbcTemplate.query(sql, args, argTypes, rowCallbackHandler);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    @Deprecated
    public void query(String dsName, String sql, Object[] args, RowCallbackHandler rowCallbackHandler) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            this.jdbcTemplate.query(sql, args, rowCallbackHandler);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    public <T> List<T> query(String dsName, PreparedStatementCreator preparedStatementCreator, RowMapper<T> rowMapper) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            return this.jdbcTemplate.query(preparedStatementCreator, rowMapper);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    public <T> List<T> query(String dsName, String sql, PreparedStatementSetter preparedStatementSetter, RowMapper<T> rowMapper) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            return this.jdbcTemplate.query(sql, preparedStatementSetter, rowMapper);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    public <T> List<T> query(String dsName, String sql, Object[] args, int[] argTypes, RowMapper<T> rowMapper) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            return this.jdbcTemplate.query(sql, args, argTypes, rowMapper);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    public <T> List<T> query(String dsName, String sql, RowMapper<T> rowMapper, Object... args) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            return this.jdbcTemplate.query(sql, rowMapper, args);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    public <T> T queryForObject(String dsName, String sql, Object[] args, int[] argTypes, RowMapper<T> rowMapper) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            return this.jdbcTemplate.queryForObject(sql, args, argTypes, rowMapper);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    public <T> T queryForObject(String dsName, String sql, RowMapper<T> rowMapper, Object... args) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            return this.jdbcTemplate.queryForObject(sql, rowMapper, args);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    public <T> T queryForObject(String dsName, String sql, Object[] args, int[] argTypes, Class<T> aClass) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            return this.jdbcTemplate.queryForObject(sql, args, argTypes, aClass);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    public <T> T queryForObject(String dsName, String sql, Class<T> aClass, Object... args) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            return this.jdbcTemplate.queryForObject(sql, aClass, args);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    public Map<String, Object> queryForMap(String dsName, String sql, Object[] args, int[] argTypes) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            return this.jdbcTemplate.queryForMap(sql, args, argTypes);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    public Map<String, Object> queryForMap(String dsName, String sql, Object... args) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            return this.jdbcTemplate.queryForMap(sql, args);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    public <T> List<T> queryForList(String dsName, String sql, Object[] args, int[] argTypes, Class<T> aClass) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            return this.jdbcTemplate.queryForList(sql, args, argTypes, aClass);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    public <T> List<T> queryForList(String dsName, String sql, Class<T> aClass, Object... args) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            return this.jdbcTemplate.queryForList(sql, aClass, args);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    public List<Map<String, Object>> queryForList(String dsName, String sql, Object[] args, int[] argTypes) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            return this.jdbcTemplate.queryForList(sql, args, argTypes);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    public List<Map<String, Object>> queryForList(String dsName, String sql, Object... args) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            return this.jdbcTemplate.queryForList(sql, args);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    public int update(String dsName, PreparedStatementCreator preparedStatementCreator) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            return this.jdbcTemplate.update(preparedStatementCreator);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    public int update(String dsName, PreparedStatementCreator preparedStatementCreator, KeyHolder keyHolder) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            return this.jdbcTemplate.update(preparedStatementCreator, keyHolder);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    public int update(String dsName, String sql, PreparedStatementSetter preparedStatementSetter) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            return this.jdbcTemplate.update(sql, preparedStatementSetter);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    public int update(String dsName, String sql, Object[] args, int[] argTypes) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            return this.jdbcTemplate.update(sql, args, argTypes);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    public int update(String dsName, String sql, Object... args) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            return this.jdbcTemplate.update(sql, args);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    public int[] batchUpdate(String dsName, String sql, BatchPreparedStatementSetter batchPreparedStatementSetter) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            return this.jdbcTemplate.batchUpdate(sql, batchPreparedStatementSetter);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    public int[] batchUpdate(String dsName, String sql, List<Object[]> batchArgs) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            return this.jdbcTemplate.batchUpdate(sql, batchArgs);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    public int[] batchUpdate(String dsName, String sql, List<Object[]> batchArgs, int[] argTypes) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            return this.jdbcTemplate.batchUpdate(sql, batchArgs, argTypes);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    public <T> int[][] batchUpdate(String dsName, String sql, Collection<T> batchArgs, int batchSize, ParameterizedPreparedStatementSetter<T> pss) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            return this.jdbcTemplate.batchUpdate(sql, batchArgs, batchSize, pss);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    public <T> T execute(String dsName, CallableStatementCreator callableStatementCreator, CallableStatementCallback<T> callableStatementCallback) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            return this.jdbcTemplate.execute(callableStatementCreator, callableStatementCallback);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    public <T> T execute(String dsName, String sql, CallableStatementCallback<T> callableStatementCallback) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            return this.jdbcTemplate.execute(sql, callableStatementCallback);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    public Map<String, Object> call(String dsName, CallableStatementCreator callableStatementCreator, List<SqlParameter> declaredParameters) throws DataAccessException {
        DynamicDataSourceContextHolder.push(determineDatasource(dsName));
        try {
            return this.jdbcTemplate.call(callableStatementCreator, declaredParameters);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

}
